"""
 Copyright (c) 2022 Intel Corporation
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
      http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""
from nncf.config.schemata.basic import ARRAY_OF_NUMBERS
from nncf.config.schemata.basic import BOOLEAN
from nncf.config.schemata.basic import NUMBER
from nncf.config.schemata.basic import STRING
from nncf.config.schemata.basic import with_attributes
from nncf.config.schemata.defaults import RB_SPARSITY_SCHEDULER
from nncf.config.schemata.defaults import SPARSITY_FREEZE_EPOCH
from nncf.config.schemata.defaults import SPARSITY_LEVEL_SETTING_MODE
from nncf.config.schemata.defaults import SPARSITY_MULTISTEP_SPARSITY_LEVELS
from nncf.config.schemata.defaults import SPARSITY_MULTISTEP_STEPS
from nncf.config.schemata.defaults import SPARSITY_SCHEDULER
from nncf.config.schemata.defaults import SPARSITY_SCHEDULER_CONCAVE
from nncf.config.schemata.defaults import SPARSITY_SCHEDULER_PATIENCE
from nncf.config.schemata.defaults import SPARSITY_SCHEDULER_POWER
from nncf.config.schemata.defaults import SPARSITY_SCHEDULER_UPDATE_PER_OPTIMIZER_STEP
from nncf.config.schemata.defaults import SPARSITY_TARGET
from nncf.config.schemata.defaults import SPARSITY_TARGET_EPOCH

SCHEDULE_OPTIONS = ['polynomial', 'exponential', 'adaptive', 'multistep', ]

COMMON_SPARSITY_PARAM_PROPERTIES = {
    "sparsity_level_setting_mode": with_attributes(STRING,
                                                   description="The mode of sparsity level setting. \n"
                                                               "`global` - the sparsity level is calculated across all "
                                                               "weight values in the network across layers, "
                                                               "`local` - the sparsity level can be set per-layer and "
                                                               "within each layer is computed with respect only to the "
                                                               "weight values within that layer.",
                                                   default=SPARSITY_LEVEL_SETTING_MODE),
    "schedule": with_attributes(STRING,
                                description=f"The type of scheduling to use for adjusting the target"
                                            f"sparsity level. Default - {RB_SPARSITY_SCHEDULER} for `rb_sparsity`, "
                                            f"{SPARSITY_SCHEDULER} otherwise",
                                enum=SCHEDULE_OPTIONS),
    "sparsity_target": with_attributes(NUMBER,
                                       description="Target sparsity level for the model, to be reached at the end of "
                                                   "the compression schedule.",
                                       default=SPARSITY_TARGET),
    "sparsity_target_epoch": with_attributes(NUMBER,
                                             description="Index of the epoch upon which the sparsity level "
                                                         "of the model is scheduled to become "
                                                         "equal to `sparsity_target`.",
                                             default=SPARSITY_TARGET_EPOCH),
    "sparsity_freeze_epoch": with_attributes(NUMBER,
                                             description="Index of the epoch upon which the sparsity mask will "
                                                         "be frozen and no longer trained.",
                                             default=SPARSITY_FREEZE_EPOCH),

    "update_per_optimizer_step": with_attributes(BOOLEAN,
                                                 description="Whether the function-based sparsity level schedulers "
                                                             "should update the sparsity level after each optimizer "
                                                             "step instead of each epoch step.",
                                                 default=SPARSITY_SCHEDULER_UPDATE_PER_OPTIMIZER_STEP),
    "steps_per_epoch": with_attributes(NUMBER,
                                       description="Number of optimizer steps in one epoch. Required to start proper "
                                                   "scheduling in the first training epoch if "
                                                   "`update_per_optimizer_step` is `true.`"),
    "multistep_steps": with_attributes(ARRAY_OF_NUMBERS,
                                       description="A list of scheduler steps at which to transition "
                                                   "to the next scheduled sparsity level (multistep "
                                                   "scheduler only).",
                                       default=SPARSITY_MULTISTEP_STEPS),
    "multistep_sparsity_levels": with_attributes(ARRAY_OF_NUMBERS,
                                                 description="Multistep scheduler only - Levels of sparsity to use at "
                                                             "each step of the scheduler as specified in the "
                                                             "`multistep_steps` attribute. The first sparsity level "
                                                             "will be applied immediately, "
                                                             "so the length of this list should be larger than the "
                                                             "length of the `multistep_steps` by one. The last "
                                                             "sparsity level will function as the ultimate sparsity "
                                                             "target, overriding the \"sparsity_target\" setting if it "
                                                             "is present.",
                                                 default=SPARSITY_MULTISTEP_SPARSITY_LEVELS),
    "patience": with_attributes(NUMBER,
                                description="A conventional patience parameter for the scheduler, "
                                            "as for any other standard scheduler. Specified in units "
                                            "of scheduler steps.",
                                default=SPARSITY_SCHEDULER_PATIENCE),
    "power": with_attributes(NUMBER,
                             description="For polynomial scheduler - determines the corresponding power value.",
                             default=SPARSITY_SCHEDULER_POWER),
    "concave": with_attributes(BOOLEAN, description="For polynomial scheduler - if `true`, then the target sparsity "
                                                    "level will be approached in concave manner, and in convex "
                                                    "manner otherwise.",
                               default=SPARSITY_SCHEDULER_CONCAVE),
}
